#include <bezier_predict/bezier_predict.h>
#include <stdio.h>
#include <ros/ros.h>
#include <ros/console.h>
#include <iostream>
#include <fstream>
#include <string>
#include <math.h>

using namespace std;
using namespace Eigen;

bool Bezierpredict::TrackingGeneration(
    const double max_vel,
    const double max_acc,
    vector<Eigen::Vector4d> predict_list)
{
  int n_seg = (int)predict_list.size();

  if (n_seg < SEGS_MIN || n_seg > SEGS_MAX) {
    cerr << "The number of predict_list is not enough or too much" << endl;
    return false;
  }

  vector<Vector3d> pos_list;
  vector<double> time_list;
  vector<double> weight_list;

  bool init_flag = false;
  double time_start;

  for (int i = 0; i < n_seg; i++) {
    pos_list.push_back(predict_list[i].head(3));
    if (!init_flag) {
      time_start = predict_list[i][3];
      init_flag = true;
    }
    time_history_ = predict_list[i][3] - time_start;
    time_list.push_back(time_history_);
  }

  for (int i = 0; i < n_seg; i++) {
    double tanh_input = time_list[n_seg - 1] - time_list[i];
    if (!tanh_input) {
      weight_list.push_back(1);
    }
    else {
      tanh_input = 1.0 / tanh_input;
      weight_list.push_back(tanh(0.6 * tanh_input));
    }
  }

  time_total_ = time_scale_ * time_history_;

  // 取得第一个end_p
  Vector3d end_p = pos_list[n_seg - 1];

  // allocate QP problem matrices and vectores
  int vars_number = TRAJ_ORDER + 1;      // 6
  int all_vars_number = 3 * vars_number; // XYZ 18

  int n_eq = 3;                             // 等式约束个数
  int n_vel_con = vars_number - 1;          // 速度约束个数 (单维度)
  int n_vel = 3 * n_vel_con;                // 速度约束个数 (三维度)
  int n_acc_con = vars_number - 2;          // 加速度约束个数 (单维度)
  int n_ineq = 3 * (n_vel_con + n_acc_con); // 不等式约束个数(vel + acc )

  Eigen::SparseMatrix<double> hessian;
  Eigen::VectorXd gradient;
  Eigen::SparseMatrix<double> linearMatrix;
  Eigen::VectorXd lowerBound = Eigen::VectorXd::Zero(n_eq + n_ineq);
  Eigen::VectorXd upperBound = Eigen::VectorXd::Zero(n_eq + n_ineq);

  Eigen::MatrixXd A = Eigen::MatrixXd::Zero(n_eq + n_ineq, all_vars_number);
  Eigen::MatrixXd A_eq = Eigen::MatrixXd::Zero(n_eq, all_vars_number);
  Eigen::VectorXd l_eq(n_eq), u_eq(n_eq);
  Eigen::MatrixXd A_ineq = Eigen::MatrixXd::Zero(n_ineq, all_vars_number);
  Eigen::VectorXd l_ineq(n_ineq), u_ineq(n_ineq);

  // 填充 A_eq, l_eq, u_eq
  for (int i = 0; i < n_eq; ++i) {
    l_eq(i) = end_p[i];
    u_eq(i) = end_p[i];
  }
  double u = time_history_ / time_total_;
  for (int i = 0; i < 3; i++) {
    for (int j = 0; j < vars_number; j++) {
      A_eq(i, i * vars_number + j) =
          C_[j] * pow(u, j) * pow((1 - u), (TRAJ_ORDER - j));
    }
  }

  // 填充 A_ineq, l_ineq, u_ineq
  // vel
  double val = TRAJ_ORDER / time_total_;
  for (int i = 0; i < 3; i++) {
    for (int j = 0; j < n_vel_con; j++) {
      l_ineq(i * n_vel_con + j) = -max_vel;
      u_ineq(i * n_vel_con + j) = max_vel;
      A_ineq(i * n_vel_con + j, i * vars_number + j) = -1.0 * val;
      A_ineq(i * n_vel_con + j, i * vars_number + j + 1) = 1.0 * val;
    }
  }
  // acc
  val = TRAJ_ORDER * (TRAJ_ORDER - 1) / pow(time_total_, 2);
  for (int i = 0; i < 3; i++) {
    for (int j = 0; j < n_acc_con; j++) {
      l_ineq(n_vel + i * n_acc_con + j) = -max_acc;
      u_ineq(n_vel + i * n_acc_con + j) = max_acc;
      A_ineq(n_vel + i * n_acc_con + j, i * vars_number + j) = 1.0 * val;
      A_ineq(n_vel + i * n_acc_con + j, i * vars_number + j + 1) = -2.0 * val;
      A_ineq(n_vel + i * n_acc_con + j, i * vars_number + j + 2) = 1.0 * val;
    }
  }

  A.block(0, 0, n_eq, all_vars_number) = A_eq;
  A.block(n_eq, 0, n_ineq, all_vars_number) = A_ineq;
  linearMatrix = A.sparseView();

  lowerBound.head(n_eq) = l_eq;
  upperBound.head(n_eq) = u_eq;
  lowerBound.tail(n_ineq) = l_ineq;
  upperBound.tail(n_ineq) = u_ineq;

  /**
   * The matrix form of Bezier curve
   *  B(t) = T(t) * M * P
   *    T(t) = [1, t, t^2, t^3, t^4, t^5]
   *    M = [1, 0, 0, 0, 0, 0,
   *       -5, 5, 0, 0, 0, 0,
   *       10, -20, 10, 0, 0, 0,
   *       -10, 30, -30, 10, 0, 0,
   *       5, -20, 30, -20, 5, 0,
   *       -1, 5, -10, 10, -5, 1]
   *    P = [P0, P1, P_square, P3, P4, P5]^T
   *
   * Why M is like this?
   *  B(t) = Sum(Bi(t) * Pi)
   *  Bi(t) = C^i_n  * t^i * (1-t)^(n-i)  t\in [0,1]
   *  such as B0(t) = C^0_5 * t^0 * (1-t)^5 = (1-t)^5
   *                = 1 - 5t + 10t^2 - 10t^3 + 5t^4 - t^5 (the first column of M)
   *          B1(t) = C^1_5 * t^1 * (1-t)^4 = 5t(1-t)^4
   *                = 5t - 20t^2 + 30t^3 - 20t^4 + 5t^5 (the second column of M)
   *
   *
   * The optimization problem is:
   *   min J = sum||T(ti)BP - pi||^2 + lambda * sum||D_n-2 * D_n-1 * P||^2
   *
   *  First item is the distance between the received point and the bezier curve
   *  first = (sum) (T(ti) * M * P - pi)^2
   *        = (sum) p^T * p - 2 * p^T * T * M * P + P^T * M^T * T^T * T * M * P
   *        linear term coefficient = -2 * p^T * T * M
   *        quadratic term coefficient =  M^T * T^T * T * M
   *
   *  Second item is the smoothness of the control point of bezier curve (acceleration)
   *  D_n-1 = n * [-1, 1, 0, 0 ...
   *               0, -1, 1, 0 ...
   *               0, 0, -1, 1 ...]  dimension = (n-1)*(n)
   *    because of the property of bezier curve, the derivative of bezier curve is also
   *    a bezier curve, can be represented by the control point and Bernstein Polynomial.
   *    So D_n-1 * P is the derivative of the bezier curve, and D_n-2 * D_n-1 * P is the
   *    second derivative of the bezier curve.
   *  second = lambda * sum||D_n-2 * D_n-1 * P||^2
   *         = lambda * P^T * D_n-1^T * D_n-2^T * D_n-2 * D_n-1 * P
   *         linear term coefficient = 0
   *         quadratic term coefficient = lambda * D_n-1^T * D_n-2^T * D_n-2 * D_n-1
   *                                    = lambda * S
   */

  // 开始定义优化对象
  MatrixXd S = MatrixXd::Zero(all_vars_number, all_vars_number); // for smooth cost
  MatrixXd D = MatrixXd::Zero(all_vars_number, all_vars_number); // for distance cost
  MatrixXd MTWTM = MatrixXd::Zero(vars_number, vars_number);
  MatrixXd TWT = MatrixXd::Zero(vars_number, vars_number);
  MatrixXd TWWT = MatrixXd::Zero(vars_number, vars_number);
  MatrixXd M = MatrixXd::Zero(all_vars_number, all_vars_number); // map matrix
  MatrixXd pT = MatrixXd::Zero(1, all_vars_number);

  // S = D_n-1^T * D_n-2^T * D_n-2 * D_n-1
  MatrixXd S_part = S_ / pow(time_total_, 4);
  // M
  MatrixXd M_part = getM(vars_number, time_total_);
  for (int j = 0; j < 3; j++) {
    S.block(j * vars_number, j * vars_number, vars_number, vars_number) = S_part;
    M.block(j * vars_number, j * vars_number, vars_number, vars_number) = M_part;
  }

  for (double j = 0; j < n_seg; j += 1) {
    // -2 * p^T * T
    pT.block(0, 0, 1, 3 * (vars_number)) +=
        getCt(time_list[j], pos_list[j]) * weight_list[j];
    // T^T * T if have weight then T^T * W * T
    MatrixXd TT = getTT(time_list[j]);
    TWT += TT * weight_list[j];
    TWWT += TT * weight_list[j] * weight_list[j];
  }

  /* linear term coefficient */
  // -2 * p^T * T * M
  gradient = (pT * M).transpose();

  /* quadratic term coefficient */
  // lambda * D_n-1^T * D_n-2^T * D_n-2 * D_n-1
  S = Lambda_ACC * S;

  // M^T * T^T * T * M  if have weight then M^T * T^T * W * T * M
  MTWTM = M_part.transpose() * TWT * M_part;
  for (int i = 0; i < 3; i++) {
    // 这个2是OOPQ要求的
    D.block(i * vars_number, i * vars_number,
            vars_number, vars_number) = 2 * MTWTM;
  }

  hessian = (S + D).sparseView();

  /* Create the osqp solver */
  OsqpEigen::Solver solver;

  // settings
  solver.settings()->setVerbosity(false); // ban the output
  solver.settings()->setWarmStart(true);
  solver.settings()->setPolish(true);

  // set the initial data of the QP solver
  solver.data()->setNumberOfVariables(all_vars_number);
  solver.data()->setNumberOfConstraints(n_eq + n_ineq);
  if (!solver.data()->setHessianMatrix(hessian))
    return false;
  if (!solver.data()->setGradient(gradient))
    return false;
  if (!solver.data()->setLinearConstraintsMatrix(linearMatrix))
    return false;
  if (!solver.data()->setLowerBound(lowerBound))
    return false;
  if (!solver.data()->setUpperBound(upperBound))
    return false;

  // instantiate the solver
  if (!solver.initSolver())
    return false;

  // solve the QP problem
  if (solver.solveProblem() != OsqpEigen::ErrorExitFlag::NoError)
    return false;

  // Record the solution
  PolyCoeff_ = solver.getSolution();
  poly_end_ = time_total_;
  poly_start_ = time_history_;
  t_start_gl_ = predict_list[0](3);
  t_predict_start_gl_ = predict_list[n_seg - 1](3);

  // Record intermediate quantities for variance prediction
  /**
   *  \hat{P} = (M'*T'*W*T*M + lambda*S)^-1 * M'*T'*W*p
   *            E = M'*T'*W*T*M + lambda*S
   *  Var(P) = Var(sigma)*E^-1*M'*T'*W*W*T*M*E^-1
   */
  MatrixXd MTWWTM = M_part.transpose() * TWWT * M_part;
  MatrixXd Quad_inv = (MTWTM + Lambda_ACC * S_part).inverse();
  Evaluation_ = M_part * Quad_inv * MTWWTM * Quad_inv * M_part.transpose();

  Eigen::VectorXd T_start = getT(poly_start_);
  prop_start_ = T_start.dot(Evaluation_ * T_start);

  // calculate the variance of parameter and adjust R^2
  /**
   * after test, when R2 < 0.99 in simulation, the traj seem changed a lot
   */
  Vector3d P_mean = Vector3d::Zero();
  VectorXd P_square;
  Vector3d P_square_sum = Vector3d::Zero();
  Vector3d SSE_w = Vector3d::Zero();
  Vector3d SST_w = Vector3d::Zero();

  for (int i = 0; i < n_seg; i++)
  {
    P_mean += pos_list[i];
    P_square = pos_list[i].array().square() * weight_list[i];
    P_square_sum += P_square;
  }
  P_mean = P_mean / n_seg;

  for (int i = 0; i < n_seg; i++)
  {
    Vector3d delta_mean2 = (pos_list[i] - P_mean).array().square();
    SST_w += weight_list[i] * delta_mean2;
  }

  for (int i = 0; i < 3; i++)
  {
    VectorXd P = PolyCoeff_.block(i * (vars_number), 0, vars_number, 1);
    VectorXd G = gradient.block(i * (vars_number), 0, vars_number, 1);

    SSE_w(i) = P.transpose() * MTWTM * P + G.dot(P) + P_square_sum(i);
  }

  Sigma2_ = SSE_w / (n_seg - vars_number);
  R2_ = 1 - SSE_w.sum() / SST_w.sum();
  R2_ = 1 - (n_seg - 1) / (n_seg - vars_number) * (1 - R2_);

  // Sigma2_ can  also be calculated by the following code,
  //   but above is more efficient

  // Vector3d Sigma2 = Vector3d::Zero();
  // Matrix3d Var = Matrix3d::Zero();

  // for (int j = 0; j < n_seg; j++)
  // {
  //   Vector3d delta = pos_list[j] - getPos(time_list[j]);
  //   Var += weight_list[j] * delta * delta.transpose();
  //   for (int i = 0; i < 3; i++)
  //   {
  //     Sigma2(i) += weight_list[j] *
  //                 pow(delta(i), 2);
  //   }
  // }
  // Sigma2 = Sigma2 / (n_seg - vars_number);

  return true;
}
